import json

from app.project.doc_to_recommendation.llm.model.TongyiChat_base_model import TongyiChatBaseModel
from app.project.doc_to_recommendation.llm.register.llm_register import LLM_REGISTER
from app.project.doc_to_recommendation.utils.generate_utils import img_to_base64
from app.project.doc_to_recommendation.utils.log_utils import get_logger

@LLM_REGISTER.register_model("tongyi_chat_img_model")
class TongyiChatImgModel(TongyiChatBaseModel):

    def __init__(self, config: dict):
        super().__init__(config)

    def bind_tools_df(self):
        pass

    def agent_calls(self, text, image, prompt=None):
        # 将image转为base64编码
        image_url = f"data:image/jpeg;base64,{img_to_base64(image['file'], image['type'])}"
        img_result = super().agent_calls(image['fpath_name'], image_url, prompt)
        logger = get_logger(__name__)
        if self.bind_t:
            if len(img_result.tool_calls) >= 1:
                imgDscp = img_result.tool_calls[0]['args']
                print("figure:", text, " type:", imgDscp['type'], " description:", imgDscp['description'])
                logger.info("figure:{} type:{} description:{}".format(text, imgDscp['type'], imgDscp['description']))
                if imgDscp['type'] == 'img':
                    return {'name': text, 'type': imgDscp['type'], 'description': imgDscp['description']}
            else:
                print("figure:", text, " type:", img_result)
                logger.info("figure:{} type:{}".format(text, img_result))
                return {'name': text, 'type': 'img', 'description': img_result.content}
        else:
            res = img_result.content.replace('json', '', 1).replace('\n', '').replace('`', '')
            res = json.loads(res)
            print("figure:", text, " type:", res['type'], " description:", res['description'])
            logger.info("figure:{} type:{} description:{}".format(text, res['type'], res['description']))
            return {'name': text, 'type': res['type'], 'description': res['description']}
